import pandas as pd
import matplotlib.pyplot as plt
class DataFeed:
    def __init__(self, df):
        self.df = df.reset_index(drop=True)  # 重置索引
        self.current_index = 0              
    def next(self):
        self.current_index += 1
        return not self.is_finished()
    def is_finished(self):
        return self.current_index >= len(self.df)
    def get_close(self):
        return self.df.loc[self.current_index, 'close']
    def get_open(self):
        return self.df.loc[self.current_index, 'open']
    def get_date(self):
        return self.df.loc[self.current_index, 'date']
class Broker:
    def __init__(self, cash, commission):
        self.init_cash = cash        # 初始资金
        self.init_comm = commission
        self.cash = cash             # 当前现金  
        self.cost = 0                # 成本
        self.spend = 0               # 花费
        self.value = 0 
        self.position = 0            # 当前持仓股数
        self.history = []            # 交易历史记录（买入/卖出）  
    def buy(self,date,price,close,size=1,):
        comm =price*size*self.init_comm
        self.cost=price*size
        self.spend=self.cost+comm
        if self.cash >= self.spend:
            self.cash -= self.spend
            self.position += size
            self.value=self.cash+close*self.position
            self.history.append(str(date)+f':buy, 价格{price}, 数量{size}, 成本{self.cost:.2f}, 佣金{comm:.2f}, 开盘买入后价值{self.value:.2f}')
    def sell(self,date,price,size=1,):
        comm =price*size*self.init_comm
        if self.position >= size:
            self.cash =self.cash+ price * size-comm
            self.position -= size
            self.value=self.cash+price*self.position
            self.history.append(str(date)+f':sell, 价格{price}, 数量{size}, 成本{self.cost:.2f}, 佣金{comm:.2f}, 收盘后净值{self.value:.2f}')
    def close(self,date,price,):
        comm =price*self.position*self.init_comm
        self.cash =self.cash + price * self.position-comm
        self.value=self.cash
        self.history.append(str(date)+f':sell, 价格{price}, 数量{self.position}, 成本{self.cost:.2f}, 佣金{comm:.2f}, 收盘后净值{self.value:.2f}')
        self.position = 0
    def get_value(self, close):
        """当前账户净值 = 现金 + 当前持仓市值"""
        return self.cash + self.position * close
class Strategy:
    params=dict(printlog=False)
    def __init__(self, data, broker,ma):
        self.data = data
        self.broker = broker
        self.values = []  # 存储每日净值
        self.order=0
        self.size_stock=0 
        self.size_stock0=0
        self.comm0=0
        self.cash0=0
        self.ma=ma-1  
    def next(self):
        """每根K线调用一次,策略逻辑写在这里"""
        date = self.data.get_date()
        price = self.data.get_open()
        close = self.data.get_close()  
        if not self.values:
            self.size_stock0=int(0.98*self.broker.cash/price)
            self.comm0=self.broker.init_comm*self.size_stock0*price
            self.cash0=self.broker.init_cash -self.size_stock0*price-self.comm0
        benchmark_value=self.size_stock0*close-self.comm0+self.cash0
        if self.order==1:
            self.broker.buy(date,price,close,size=self.size_stock)
            self.order=0     
        elif self.order==2:
            self.broker.close(date,price,)
            self.order=0
        self.size_stock=0  
        # 价格高于过去n天均值就买入，低于就卖出
        if self.data.current_index >= self.ma:
            prev = self.data.df.loc[self.data.current_index - self.ma: self.data.current_index , 'close']
            ma = prev.mean() 
            indicator=round(close - ma,5)                
            if not self.broker.position  and indicator > 0:
                self.order=1
                self.size_stock=int(0.98*self.broker.cash/close)
            elif indicator < 0 and self.broker.position:
                self.order=2
            else:pass
        self.values.append([date, self.broker.get_value(close),benchmark_value])
# 回测主函数
def run_backtest(csv_file,ma=30,cash=10000,commission=0.0002):
    df = pd.read_csv(csv_file)  # 加载CSV数据，要求有 date, close，open 字段
    data = DataFeed(df)         
    broker = Broker(cash,commission)      
    strategy = Strategy(data, broker,ma)  
    while not data.is_finished():
        strategy.next()
        data.next()
    # 打印交易历史到backlog.txt
    with open('backlog.txt', 'w', encoding='utf-8') as f:
        f.write("")
        f.close()
    for _ in broker.history:
        with open('backlog.txt', 'a', encoding='utf-8') as f:
                f.write(str(_)+'\n')
    dates, values ,benchmark_values= zip(*strategy.values)
    print(f"最终总资金: {values[-1]:.2f}")
    plt.rcParams["font.sans-serif"] = ["SimHei"]
    plt.rcParams["axes.unicode_minus"] = False
    plt.figure(figsize=(8,5))
    if len(dates)<1000:
        plt.plot(dates,values,label='策略收益')
        plt.plot(dates,benchmark_values,label='基准收益')
    else:
        plt.plot(values,label='策略收益',)
        plt.plot(benchmark_values,label='基准收益')
    plt.legend()
    plt.title('资产净值')
    plt.xlabel(f'Date:{dates[0]}-----{dates[-1]}')
    plt.ylabel('总资金')
    plt.grid()
    plt.xticks(rotation=45)
    plt.tight_layout()
    plt.show()
if __name__ == '__main__':
    csv_file='./股指/纳斯达克-19910812-20250606.csv'
    cash=1000000
    commission=0.0002
    ma=30
    run_backtest(csv_file,ma,cash,commission)